(cl:in-package #:sicl-register-allocation)

(defun find-register-in-arrangement (arrangement datum &key (accept-stack nil))
  ;; Any datum that is not a lexical location does not need to be
  ;; replaced.
  (unless (typep datum '(or cleavir-ir:lexical-location
                            cleavir-ir:raw-datum))
    (return-from find-register-in-arrangement datum))
  (multiple-value-bind (stack-slot register-number)
      (arr:find-attribution arrangement datum)
    (cond
      ((not (null register-number))
       (aref x86-64:*registers* register-number))
      ;; Try to make a useful error...
      ((not (null stack-slot))
       (if accept-stack
           (make-instance 'cleavir-ir:stack-location
                          :offset stack-slot)
           (error "~S has an attributed stack slot but no attributed register.
Did you forget to call ENSURE-INPUT-AVAILABLE?"
                  datum)))
      (t
       (error "~S has no attribution." datum)))))

;; A set of instructions which were generated as part of introducing
;; registers, and thus should not be processed.
(defvar *generated-instructions* (make-hash-table))

(defun mark-as-generated (instruction)
  (setf (gethash instruction *generated-instructions*) t)
  instruction)

(defun instruction-generated-p (instruction)
  (values
   (gethash instruction *generated-instructions* nil)))

(defun save-to-stack-instruction (from-register to-stack-slot)
  (mark-as-generated
   (x86-64:save-to-stack-instruction
    (aref x86-64:*registers* from-register)
    to-stack-slot)))

(defun load-from-stack-instruction (from-stack-slot to-register)
  (mark-as-generated
   (x86-64:load-from-stack-instruction
    from-stack-slot
    (aref x86-64:*registers* to-register))))

(defun replace-instruction (instruction new-instruction)
  (cleavir-ir:insert-instruction-after new-instruction instruction)
  (cleavir-ir:delete-instruction instruction)
  new-instruction)

(defgeneric introduce-registers-for-instruction (instruction))

;;; Skip instructions which were generated by processing previous instructions.
(defmethod introduce-registers-for-instruction :around (instruction)
  (unless (instruction-generated-p instruction)
    (call-next-method)))

;;; For most instructions, it suffices to replace every lexical
;;; location with the registers attributed to them.
(defmethod introduce-registers-for-instruction
    ((instruction cleavir-ir:instruction))
  (let ((input-arrangement (input-arrangement instruction)))
    (setf (cleavir-ir:inputs instruction)
          (mapcar (lambda (input)
                    (find-register-in-arrangement input-arrangement
                                                  input))
                  (cleavir-ir:inputs instruction))))
  (let ((output-arrangement (output-arrangement instruction)))
    (setf (cleavir-ir:outputs instruction)
          (mapcar (lambda (output)
                    (find-register-in-arrangement output-arrangement
                                                  output))
                  (cleavir-ir:outputs instruction)))))

;;; An ENTER-INSTRUCTION only has outputs, and only the first two
;;; outputs (static and dynamic environment locations) are attributed.

(defmethod introduce-registers-for-instruction
    ((instruction cleavir-ir:enter-instruction))
  (let ((output-arrangement (output-arrangement instruction))
        (outputs (cleavir-ir:outputs instruction)))
    (setf (cleavir-ir:outputs instruction)
          (list* (find-register-in-arrangement output-arrangement
                                               (first outputs))
                 (find-register-in-arrangement output-arrangement
                                               (second outputs))
                 (nthcdr 2 outputs)))))

;;; Assignment instructions generated by SPILL and UNSPILL need to
;;; be converted to MEMSET2 and MEMREF2 instructions though.

(defmethod introduce-registers-for-instruction
    ((instruction cleavir-ir:assignment-instruction))
  (let ((input  (first (cleavir-ir:inputs instruction)))
        (output (first (cleavir-ir:outputs instruction))))
    ;; Such assignments assign a lexical location to itself, so
    ;; any other assignments can be handled with the normal rule.
    (unless (and (eq input output)
                 (typep input 'cleavir-ir:lexical-location))
      (return-from introduce-registers-for-instruction
        (call-next-method)))
    (multiple-value-bind (in-stack in-register)
        (arr:find-attribution (input-arrangement instruction) input)
      (assert (not (and (null in-register)
                        (null in-stack)))
              ()
              "~S has no attribution in the input arrangement of ~S"
              input instruction)
      (multiple-value-bind (out-stack out-register)
          (arr:find-attribution (output-arrangement instruction) output)
        (assert (not (and (null out-register)
                          (null out-stack)))
              ()
              "~S has no attribution in the output arrangement of ~S"
              output instruction)
        (cond
          ((and (null in-stack) (not (null out-stack)))
           ;; This instruction is a spill.
           (replace-instruction
            instruction
            (save-to-stack-instruction in-register out-stack)))
          ((and (null in-register) (not (null out-register)))
           ;; This instruction is an unspill.
           (replace-instruction
            instruction
            (load-from-stack-instruction in-stack out-register)))
          ((and (not (null in-register)) (not (null in-register)))
           ;; This instruction is a plain register to register assignment.
           (setf (cleavir-ir:inputs instruction)
                 (list (aref x86-64:*registers* in-register))
                 (cleavir-ir:outputs instruction)
                 (list (aref x86-64:*registers* out-register))))
          (t
           (error "Not sure what the assignment of (~A, ~A) to (~A, ~A) is meant to be."
                  in-register in-stack out-register out-stack)))))))

;;; NOP-INSTRUCTIONs shouldn't have inputs and outputs, but we
;;; generate NOP-INSTRUCTIONs with inputs and outputs by CHANGE-CLASS
;;; somewhere.  So just ignore them until those are sniffed out.

(defmethod introduce-registers-for-instruction
    ((instruction cleavir-ir:nop-instruction))
  nil)

(defmethod introduce-registers-for-instruction
    ((instruction cleavir-ir:unreachable-instruction))
  nil)

;;; We have to do something for FUNCALL-INSTRUCTION,
;;; NAMED-CALL-INSTRUCTION and MULTIPLE-VALUE-CALL-INSTRUCTION, but I
;;; am not sure of what just yet.

(macrolet ((def (class)
             `(defmethod introduce-registers-for-instruction
                  ((instruction ,class))
                (let ((input-arrangement (input-arrangement instruction)))
                  (setf (cleavir-ir:inputs instruction)
                        (mapcar (lambda (location)
                                  (find-register-in-arrangement input-arrangement
                                                                location
                                                                :accept-stack t))
                                (cleavir-ir:inputs instruction)))
                  (setf (cleavir-ir:dynamic-environment-location instruction)
                        (find-register-in-arrangement
                         input-arrangement
                         (cleavir-ir:dynamic-environment-location instruction)
                         :accept-stack t))))))
  (def cleavir-ir:funcall-instruction)
  (def cleavir-ir:named-call-instruction)
  (def cleavir-ir:catch-instruction)
  (def cleavir-ir:bind-instruction)
  (def cleavir-ir:unwind-instruction)
  (def cleavir-ir:initialize-values-instruction)
  (def cleavir-ir:initialize-closure-instruction)
  (def cleavir-ir:enclose-instruction)
  (def cleavir-ir:multiple-value-call-instruction)
  (def sicl-ir:patch-literal-instruction))

(defun introduce-registers (mir)
  (cleavir-ir:map-local-instructions
   #'introduce-registers-for-instruction
   mir))
